import os.path
import numpy as np
from tensorflow.python.platform import gfile
import tensorflow as tf
from python_ai.common.xcommon import *

BOTTLENECK_TENSOR_NAME = 'pool_3/_reshape:0'
JPEG_DATA_TENSOR_NAME = 'DecodeJpeg/contents:0'
MODEL_DIR = r'../../../../../large_data/model/inceptionV3/'
MODEL_FILE = 'tensorflow_inception_graph.pb'
path = r'../../../../../large_data/CV3/_many_files/agriculture/train'
INPUT_IMAGE = path
OUTPUT_VEC = '../../../../../large_data/CV3/_many_files/agriculture/train_bottleneck'
os.makedirs(OUTPUT_VEC, exist_ok=True)


def run_bottleneck_on_image(sess, image_data, image_data_tensor, bottleneck_tensor):
    bottleneck_values = sess.run(bottleneck_tensor, {image_data_tensor: image_data})
    bottleneck_values = np.squeeze(bottleneck_values)
    return bottleneck_values


def get_random_cached_bottlenecks(sess, path, jpeg_data_tensor, bottleneck_tensor):
    for _, class_name in enumerate(os.listdir(path)):
        sep(class_name)
        sub_path = os.path.join(path, class_name)
        cnt = 0
        for img in os.listdir(sub_path):
            img_path=os.path.join(sub_path,img)
            image_data = gfile.FastGFile(img_path, 'rb').read()
            bottleneck_values = run_bottleneck_on_image(sess, image_data, jpeg_data_tensor, bottleneck_tensor)

            sub_dir_path = os.path.join(OUTPUT_VEC, class_name)
            if not os.path.exists(sub_dir_path):
                os.makedirs(sub_dir_path)

            new_image_path=os.path.join(sub_dir_path, img)+".txt"
            if not os.path.exists(new_image_path):
                bottleneck_string = ','.join(str(x) for x in bottleneck_values)
                with open(new_image_path, 'w') as bottleneck_file:
                    bottleneck_file.write(bottleneck_string)
            else:
                break

            cnt += 1
            if cnt % 25 == 0:
                print(f'{cnt} pictures processed.')
        if cnt % 25 != 0:
            print(f'{cnt} pictures processed.')


if __name__ == '__main__':
    with gfile.FastGFile(os.path.join(MODEL_DIR, MODEL_FILE), 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        jpeg_data_tensor, bottleneck_tensor = tf.import_graph_def(
            graph_def, return_elements=[JPEG_DATA_TENSOR_NAME, BOTTLENECK_TENSOR_NAME])
    with tf.Session() as sess:
        tf.global_variables_initializer().run()
        get_random_cached_bottlenecks(sess, INPUT_IMAGE, jpeg_data_tensor, bottleneck_tensor)
